import os
import torch
import numpy as np
def calculate_iou(box1, box2):
    # Assumes boxes are in (x_min, y_min, x_max, y_max) format
    x_left = max(box1[0], box2[0])
    y_top = max(box1[1], box2[1])
    x_right = min(box1[2], box2[2])
    y_bottom = min(box1[3], box2[3])

    # Check for no overlap condition
    if x_right <= x_left or y_bottom <= y_top:
        return 0.0

    # Calculate intersection area
    intersection_area = (x_right - x_left) * (y_bottom - y_top)

    # Calculate each box's area
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])

    # Calculate IoU
    iou = intersection_area / float(box1_area + box2_area - intersection_area)

    # Handle edge cases based on definition
    if box1_area <= 0 or box2_area <= 0:
        return 0.0  # Avoid dividing by zero or negative areas

    return iou

def compute_ex_co_co_metric(data_samples):
    # 处理每一个结果
    score_threshold = 0.5
    iou_threshold = 0.5
    metrics_details = []
    for data_sample in data_samples:
        # print(data_sample)
        img_id = data_sample["img_id"]
        pred_bboxes = data_sample["pred_instances"]["bboxes"]
        pred_scores = data_sample["pred_instances"]["scores"]
        pred_labels = data_sample["pred_instances"]["labels"]
        gt_bboxes = data_sample["gt_instances"]["bboxes"]
        gt_labels = data_sample["gt_instances"]["labels"]

        matched_gt = set()
        matched_pred = set()

        tp = 0
        fp = 0
        fn = 0

        # Calculate TP and matched predictions
        for gt_idx, gt_bbox in enumerate(gt_bboxes):
            gt_label = gt_labels[gt_idx]
            best_iou = 0
            best_pred_idx = -1

            for pred_idx, (pred_bbox, pred_score) in enumerate(zip(pred_bboxes, pred_scores)):
                if pred_idx in matched_pred or pred_score < score_threshold:
                    continue

                pred_label = pred_labels[pred_idx]

                if gt_label != pred_label:
                    continue

                iou = calculate_iou(gt_bbox, pred_bbox)
                
                if iou > best_iou:
                    best_iou = iou
                    best_pred_idx = pred_idx

            if best_iou >= iou_threshold and best_pred_idx != -1:
                tp += 1
                matched_gt.add(gt_idx)
                matched_pred.add(best_pred_idx)

        # Calculate FP
        for pred_idx, pred_score in enumerate(pred_scores):
            if pred_idx not in matched_pred and pred_score >= score_threshold:
                fp += 1

        # Calculate FN
        for gt_idx in range(len(gt_bboxes)):
            if gt_idx not in matched_gt:
                fn += 1

        # Append results for this image
        metrics_details.append({
            "img_id": img_id,
            "TP": tp,
            "FP": fp,
            "FN": fn
        })
    
    # print(metrics_details)

def polygon_masks_to_serializable(polygon_masks):
    """将 PolygonMasks 对象转换为可序列化的格式"""
    # 提取掩码数据
    serializable_masks = []

    for polygons in polygon_masks.masks:
        serializable_polygon = [polygon.tolist() for polygon in polygons]
        serializable_masks.append(serializable_polygon)

    result = {
        'masks': serializable_masks,
        'height': polygon_masks.height,
        'width': polygon_masks.width
    }
    return result

def tensor_to_list(data):
    """将 Tensor、ndarray 和 tuple 转换为可序列化的类型。"""
    if isinstance(data, torch.Tensor):
        return data.tolist()
    elif isinstance(data, np.ndarray):
        return data.tolist()
    elif isinstance(data, list):
        return [tensor_to_list(item) for item in data]
    elif isinstance(data, tuple):
        return [tensor_to_list(item) for item in data]
    elif isinstance(data, dict):
        return {key: tensor_to_list(value) for key, value in data.items() if key != 'masks'}
    elif type(data).__name__ == 'PolygonMasks':
        return polygon_masks_to_serializable(data)

    return data
